'''
pre-registration
'''
from dataprocessor.itkdatareader import SitkDataReader
import SimpleITK as sitk
import os
from NIIVisualization.Niiplot import multi_slice_viewer
from NIIVisualization.Niiplot import plot_3d
class Registrator():
    def __init__(self,image_fix,label_fix,image_moving,lable_moving):
        self.fix_imgs = SitkDataReader(image_fix)
        self.fix_labels = SitkDataReader(label_fix)
        self.mv_imgs = SitkDataReader(image_moving)
        self.mv_labels = SitkDataReader(lable_moving)

    def pre_registraion(self):
        for i in range(self.fix_imgs.num_data):
            print("processing %d"%(i))
            fix_img=self.fix_imgs.get_file_obj(i)
            fix_lab=self.fix_labels.get_file_obj(i)
            mv_img=self.mv_imgs.get_file_obj(i)
            mv_lab=self.mv_labels.get_file_obj(i)

            initial_transform = sitk.CenteredTransformInitializer(fix_img,
                                                                  mv_img,
                                                                  sitk.Euler3DTransform(),
                                                                  sitk.CenteredTransformInitializerFilter.GEOMETRY)
            mv_label_resampled = sitk.Resample(mv_lab, fix_img, initial_transform, sitk.sitkNearestNeighbor, 0.0,
                                               mv_lab.GetPixelID())


            # multi_slice_viewer(sitk.GetArrayViewFromImage(mv_label_resampled))
            # multi_slice_viewer(sitk.GetArrayViewFromImage(fix_lab))

            mv_img_resampled = sitk.Resample(mv_img, fix_img, initial_transform, sitk.sitkBSpline, 0.0, mv_img.GetPixelID())


            # mv_img_resampled3 = sitk.Resample(mv_img, initial_transform, sitk.sitkBSpline, 0.0,
            #                                   mv_img.GetPixelID())
            # multi_slice_viewer(mv_img)
            # multi_slice_viewer(fix_img)
            # multi_slice_viewer(mv_img_resampled)
            # multi_slice_viewer(mv_img_resampled)
            # multi_slice_viewer(mv_img_resampled2)
            # multi_slice_viewer(mv_img_resampled3)
            self.__itksave(fix_img,fix_lab,mv_img_resampled,mv_label_resampled,i)
    def __itksave(self,fix_img,fix_lab,mv_img,mv_lab,i):
        fix_img_path = os.path.join(self.fix_imgs.dir_name + "_pre_reg", os.path.basename(self.fix_imgs.files[i]))
        fix_lab_path = os.path.join(self.fix_labels.dir_name + "_pre_reg", os.path.basename(self.fix_labels.files[i]))
        mv_img_path = os.path.join(self.mv_imgs.dir_name + "_pre_reg", os.path.basename(self.mv_imgs.files[i]))
        mv_lab_path = os.path.join(self.mv_labels.dir_name + "_pre_reg", os.path.basename(self.mv_labels.files[i]))
        if not os.path.exists(os.path.dirname(fix_img_path)):
            os.makedirs(os.path.dirname(fix_img_path))
        if not os.path.exists(os.path.dirname(fix_lab_path)):
            os.makedirs(os.path.dirname(fix_lab_path))
        if not os.path.exists(os.path.dirname(mv_img_path)):
            os.makedirs(os.path.dirname(mv_img_path))
        if not os.path.exists(os.path.dirname(mv_lab_path)):
            os.makedirs(os.path.dirname(mv_lab_path))
        sitk.WriteImage(fix_img,fix_img_path)
        sitk.WriteImage(fix_lab,fix_lab_path)
        sitk.WriteImage(mv_img,mv_img_path)
        sitk.WriteImage(mv_lab,mv_lab_path)

if __name__=="__main__":
    reg=Registrator("E:\MIA_CODE_DATA\zhuang_data\MMWHS\CT\\train\ct-image_result",
                    "E:\MIA_CODE_DATA\zhuang_data\MMWHS\CT\\train\ct-label_result",
                    "E:\MIA_CODE_DATA\zhuang_data\MMWHS\MRI\mr-image_result",
                    "E:\MIA_CODE_DATA\zhuang_data\MMWHS\MRI\mr-label_result")
    reg.pre_registraion()
